#include "matrix.hpp"
// #include <cassert>
#include <chrono>
#include <cstdlib>
#include <iostream>
#include <mpi.h>
using namespace std;
using namespace std::chrono;
// scatter 和 gather
// N * K 矩阵 multiply K vector
void MatrixMultiplyVector(const Matrix &a, const Matrix &b, Matrix &c, int row);
int main(int argc, char *argv[]) {
  int my_rank, num_process;

  if (argc != 3) {
    cerr << "./a.out N K\n" << endl;
    exit(EXIT_FAILURE);
  }

  MPI_Init(&argc, &argv);
  MPI_Comm_size(MPI_COMM_WORLD, &num_process);
  MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

  const int N = atoi(argv[1]);
  const int K = atoi(argv[2]);
  if (my_rank == 0) {
    printf(" argc =  %d\n", argc);
    printf("\n %d*%d multiply %d\n", N, K, K);
  }

  int rows_per_process = N / num_process;

  if (my_rank == 0) {
    Matrix a(N, K), b(K, 1), c(N, 1);
    a.RandomInitialize();
    b.RandomInitialize();
    // c = 0.0;
    vector<int> displacement(num_process), send_counts(num_process);
    for (int i = 0; i < num_process; ++i) {
      send_counts[i] = rows_per_process * K;
      //   相对于起始的偏移
      displacement[i] = i * rows_per_process * K;
    }
    // 最后一个数据量可能会大一些
    send_counts[num_process - 1] =
        (N - rows_per_process * (num_process - 1)) * K;
    MPI_Scatterv(a.Data(), send_counts.data(), displacement.data(), MPI_DOUBLE,
                 MPI_IN_PLACE, 0, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(b.Data(), K, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MatrixMultiplyVector(a, b, c, rows_per_process);

    for (int i = 0; i < num_process; ++i) {
      send_counts[i] = rows_per_process;
      //   相对于起始的偏移
      displacement[i] = i * rows_per_process;
    }
    send_counts[num_process - 1] = N - rows_per_process * (num_process - 1);
    printf("hello\n");
    // MPI_Gatherv(c.Data(), rows_per_process, MPI_DOUBLE, d.Data(),
    //             send_counts.data(), displacement.data(), MPI_DOUBLE, 0,
    //             MPI_COMM_WORLD);

    // mpi_in_place 不能改为 c.data  不能相同
    MPI_Gatherv(MPI_IN_PLACE, rows_per_process, MPI_DOUBLE, c.Data(),
                send_counts.data(), displacement.data(), MPI_DOUBLE, 0,
                MPI_COMM_WORLD);
    // check
    Matrix cc(N, 1);
    MatrixMultiply(a, b, cc);
    int bool_result = cc.IsSame(c);
    printf("%d\n", bool_result);
  } else {
    if (my_rank == num_process - 1) {
      rows_per_process = N - (num_process - 1) * rows_per_process;
    }
    Matrix subA(rows_per_process, K), b(K, 1), subC(rows_per_process, 1);

    MPI_Scatterv(nullptr, nullptr, nullptr, MPI_DOUBLE, subA.Data(),
                 rows_per_process * K, MPI_DOUBLE, 0, MPI_COMM_WORLD);
    MPI_Bcast(b.Data(), K, MPI_DOUBLE, 0, MPI_COMM_WORLD);

    MatrixMultiplyVector(subA, b, subC, rows_per_process);

    MPI_Gatherv(subC.Data(), rows_per_process, MPI_DOUBLE, nullptr, nullptr,
                nullptr, MPI_DOUBLE, 0, MPI_COMM_WORLD);
  }

  MPI_Finalize();
  return 0;
}

void MatrixMultiplyVector(const Matrix &a, const Matrix &b, Matrix &c,
                          int row) {
  // assert((a.GetColumn() == b.()) && (row <= c.size()));
  int column = a.GetColumn();
  for (int i = 0; i < row; ++i) {
    double tmp{0.0};
    for (int j = 0; j < column; ++j) {
      tmp += a(i, j) * b(j, 0);
    }
    c(i, 0) = tmp;
  }
}